package com.wt.admin.util;

import lombok.Getter;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadPoolExecutor;

public class ThreadPoolUtil {

    @Getter
    private final ThreadPoolTaskExecutor executor;
    // 非核心线程存活时间
    private final Integer keepSeconds = 30;

    public ThreadPoolUtil(int core, int max, int queueSize, String name) {this(core,max,queueSize, new ThreadPoolExecutor.DiscardPolicy(),name);}
    public ThreadPoolUtil(int core, int max, int queueSize) {
        this(core,max,queueSize,null);
    }
    public ThreadPoolUtil(int core, int max, int queueSize, RejectedExecutionHandler rejected, String name) {
        this.executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(core);
        executor.setMaxPoolSize(max);
        executor.setQueueCapacity(queueSize);
        executor.setKeepAliveSeconds(keepSeconds);
        executor.setWaitForTasksToCompleteOnShutdown(true);
        executor.setThreadNamePrefix(name+"-");
        executor.setRejectedExecutionHandler(rejected);
        executor.initialize();
    }


    public void execute(Runnable runnable){
        executor.execute(runnable);
    }

    public void shutdown() {
        executor.shutdown();
    }

    public static class HashBasedThreadPool {
        private final ThreadPoolUtil[] threadPools;
        private final int poolCount;

        public HashBasedThreadPool(int poolCount,String name) {
            this.poolCount = poolCount;
            this.threadPools = new ThreadPoolUtil[poolCount];

            for (int i = 0; i < poolCount; i++) {
                threadPools[i] = new ThreadPoolUtil(1,1,50,name+"-HASH-"+i);
            }
        }

        public void execute(Runnable task, Object key) {
            int hash = key.hashCode() % poolCount;
            threadPools[Math.abs(hash)].execute(task);
        }

        public void shutdown() {
            for (ThreadPoolUtil pool : threadPools) {
                pool.shutdown();
            }
        }
    }

}
